import sys
import types
from opcode import *
from opcode import __all__ as _opcodes_all
__all__ = (['dis',
 'disassemble',
 'distb',
 'disco',
 'findlinestarts',
 'findlabels'] + _opcodes_all)
del _opcodes_all
_have_code = (types.MethodType,
 types.FunctionType,
 types.CodeType,
 types.ClassType,
 type)

def dis(x = None):
    if (x is None):
        distb()
        return 
    if isinstance(x, types.InstanceType):
        x = x.__class__
    if hasattr(x, 'im_func'):
        x = x.im_func
    if hasattr(x, 'func_code'):
        x = x.func_code
    if hasattr(x, '__dict__'):
        items = x.__dict__.items()
        items.sort()
        for (name, x1,) in items:
            if isinstance(x1, _have_code):
                print ('Disassembly of %s:' % name)
                try:
                    dis(x1)
                except TypeError, msg:
                    print 'Sorry:',
                    print msg
                print

    elif hasattr(x, 'co_code'):
        disassemble(x)
    elif isinstance(x, str):
        disassemble_string(x)
    else:
        raise TypeError, ("don't know how to disassemble %s objects" % type(x).__name__)



def distb(tb = None):
    if (tb is None):
        try:
            tb = sys.last_traceback
        except AttributeError:
            raise RuntimeError, 'no last traceback to disassemble'
        while tb.tb_next:
            tb = tb.tb_next

    disassemble(tb.tb_frame.f_code, tb.tb_lasti)



def disassemble(co, lasti = -1):
    code = co.co_code
    labels = findlabels(code)
    linestarts = dict(findlinestarts(co))
    n = len(code)
    i = 0
    extended_arg = 0
    free = None
    while (i < n):
        c = code[i]
        op = ord(c)
        if (i in linestarts):
            if (i > 0):
                print
            print ('%3d' % linestarts[i]),
        else:
            print '   ',
        if (i == lasti):
            print '-->',
        else:
            print '   ',
        if (i in labels):
            print '>>',
        else:
            print '  ',
        print repr(i).rjust(4),
        print opname[op].ljust(20),
        i = (i + 1)
        if (op >= HAVE_ARGUMENT):
            oparg = ((ord(code[i]) + (ord(code[(i + 1)]) * 256)) + extended_arg)
            extended_arg = 0
            i = (i + 2)
            if (op == EXTENDED_ARG):
                extended_arg = (oparg * 65536L)
            print repr(oparg).rjust(5),
            if (op in hasconst):
                print (('(' + repr(co.co_consts[oparg])) + ')'),
            elif (op in hasname):
                print (('(' + co.co_names[oparg]) + ')'),
            else:
                if (op in hasjrel):
                    print (('(to ' + repr((i + oparg))) + ')'),
            if (op in haslocal):
                print (('(' + co.co_varnames[oparg]) + ')'),
            elif (op in hascompare):
                print (('(' + cmp_op[oparg]) + ')'),
            else:
                if (op in hasfree):
                    if (free is None):
                        free = (co.co_cellvars + co.co_freevars)
                    print (('(' + free[oparg]) + ')'),
        print




def disassemble_string(code, lasti = -1, varnames = None, names = None, constants = None):
    labels = findlabels(code)
    n = len(code)
    i = 0
    while (i < n):
        c = code[i]
        op = ord(c)
        if (i == lasti):
            print '-->',
        else:
            print '   ',
        if (i in labels):
            print '>>',
        else:
            print '  ',
        print repr(i).rjust(4),
        print opname[op].ljust(15),
        i = (i + 1)
        if (op >= HAVE_ARGUMENT):
            oparg = (ord(code[i]) + (ord(code[(i + 1)]) * 256))
            i = (i + 2)
            print repr(oparg).rjust(5),
            if (op in hasconst):
                if constants:
                    print (('(' + repr(constants[oparg])) + ')'),
                else:
                    print ('(%d)' % oparg),
            else:
                if (op in hasname):
                    if (names is not None):
                        print (('(' + names[oparg]) + ')'),
                    else:
                        print ('(%d)' % oparg),
            if (op in hasjrel):
                print (('(to ' + repr((i + oparg))) + ')'),
            elif (op in haslocal):
                if varnames:
                    print (('(' + varnames[oparg]) + ')'),
                else:
                    print ('(%d)' % oparg),
            else:
                if (op in hascompare):
                    print (('(' + cmp_op[oparg]) + ')'),
        print



disco = disassemble

def findlabels(code):
    labels = []
    n = len(code)
    i = 0
    while (i < n):
        c = code[i]
        op = ord(c)
        i = (i + 1)
        if (op >= HAVE_ARGUMENT):
            oparg = (ord(code[i]) + (ord(code[(i + 1)]) * 256))
            i = (i + 2)
            label = -1
            if (op in hasjrel):
                label = (i + oparg)
            elif (op in hasjabs):
                label = oparg
            if (label >= 0):
                if (label not in labels):
                    labels.append(label)

    return labels



def findlinestarts(code):
    byte_increments = [ ord(c) for c in code.co_lnotab[0None2] ]
    line_increments = [ ord(c) for c in code.co_lnotab[1None2] ]
    lastlineno = None
    lineno = code.co_firstlineno
    addr = 0
    for (byte_incr, line_incr,) in zip(byte_increments, line_increments):
        if byte_incr:
            if (lineno != lastlineno):
                yield (addr, lineno)
                lastlineno = lineno
            addr += byte_incr
        lineno += line_incr

    if (lineno != lastlineno):
        yield (addr, lineno)



def _test():
    if sys.argv[1:]:
        if sys.argv[2:]:
            sys.stderr.write('usage: python dis.py [-|file]\n')
            sys.exit(2)
        fn = sys.argv[1]
        if ((not fn) or (fn == '-')):
            fn = None
    else:
        fn = None
    if (fn is None):
        f = sys.stdin
    else:
        f = open(fn)
    source = f.read()
    if (fn is not None):
        f.close()
    else:
        fn = '<stdin>'
    code = compile(source, fn, 'exec')
    dis(code)


if (__name__ == '__main__'):
    _test()

